package com.harmony.pilot.knowledge.service;

import com.google.gson.Gson;
import com.google.gson.JsonObject;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * 向量嵌入服务
 * 将代码文本转换为向量，用于语义搜索
 */
@Slf4j
@Service
public class EmbeddingService {
    
    @Value("${knowledge.embedding.provider:GLM}")
    private String provider;
    
    @Value("${knowledge.embedding.model:glm-4-embedding}")
    private String model;
    
    @Value("${knowledge.embedding.api-key:}")
    private String apiKey;
    
    @Value("${knowledge.embedding.dimension:1024}")
    private int dimension;
    
    private final OkHttpClient httpClient;
    private final Gson gson;
    
    private static final String GLM_API_URL = "https://open.bigmodel.cn/api/paas/v4/embeddings";
    private static final String OPENAI_API_URL = "https://api.openai.com/v1/embeddings";
    
    public EmbeddingService() {
        this.httpClient = new OkHttpClient.Builder()
            .connectTimeout(30, TimeUnit.SECONDS)
            .readTimeout(60, TimeUnit.SECONDS)
            .build();
        this.gson = new Gson();
    }
    
    /**
     * 生成单个文本的embedding
     */
    public float[] generateEmbedding(String text) throws IOException {
        List<String> texts = new ArrayList<>();
        texts.add(text);
        
        List<float[]> embeddings = generateEmbeddings(texts);
        return embeddings.isEmpty() ? new float[0] : embeddings.get(0);
    }
    
    /**
     * 批量生成embeddings
     */
    public List<float[]> generateEmbeddings(List<String> texts) throws IOException {
        log.info("生成embeddings: provider={}, count={}", provider, texts.size());
        
        if ("GLM".equalsIgnoreCase(provider)) {
            return generateGLMEmbeddings(texts);
        } else if ("OpenAI".equalsIgnoreCase(provider)) {
            return generateOpenAIEmbeddings(texts);
        } else {
            log.warn("未知的embedding provider: {}, 使用默认GLM", provider);
            return generateGLMEmbeddings(texts);
        }
    }
    
    /**
     * 使用GLM生成embeddings
     */
    private List<float[]> generateGLMEmbeddings(List<String> texts) throws IOException {
        // 构建请求
        JsonObject requestBody = new JsonObject();
        requestBody.addProperty("model", model);
        requestBody.add("input", gson.toJsonTree(texts));
        
        Request request = new Request.Builder()
            .url(GLM_API_URL)
            .addHeader("Authorization", "Bearer " + apiKey)
            .addHeader("Content-Type", "application/json")
            .post(RequestBody.create(
                gson.toJson(requestBody),
                MediaType.parse("application/json")
            ))
            .build();
        
        // 发送请求
        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("GLM API调用失败: " + response.code() + " - " + response.message());
            }
            
            String responseBody = response.body().string();
            JsonObject jsonResponse = gson.fromJson(responseBody, JsonObject.class);
            
            // 解析embeddings
            List<float[]> embeddings = new ArrayList<>();
            var dataArray = jsonResponse.getAsJsonArray("data");
            
            for (int i = 0; i < dataArray.size(); i++) {
                var embeddingObj = dataArray.get(i).getAsJsonObject();
                var embeddingArray = embeddingObj.getAsJsonArray("embedding");
                
                float[] embedding = new float[embeddingArray.size()];
                for (int j = 0; j < embeddingArray.size(); j++) {
                    embedding[j] = embeddingArray.get(j).getAsFloat();
                }
                embeddings.add(embedding);
            }
            
            log.info("成功生成 {} 个embeddings, 维度: {}", embeddings.size(), 
                    embeddings.isEmpty() ? 0 : embeddings.get(0).length);
            
            return embeddings;
        }
    }
    
    /**
     * 使用OpenAI生成embeddings
     */
    private List<float[]> generateOpenAIEmbeddings(List<String> texts) throws IOException {
        // 构建请求
        JsonObject requestBody = new JsonObject();
        requestBody.addProperty("model", "text-embedding-ada-002");
        requestBody.add("input", gson.toJsonTree(texts));
        
        Request request = new Request.Builder()
            .url(OPENAI_API_URL)
            .addHeader("Authorization", "Bearer " + apiKey)
            .addHeader("Content-Type", "application/json")
            .post(RequestBody.create(
                gson.toJson(requestBody),
                MediaType.parse("application/json")
            ))
            .build();
        
        // 发送请求
        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("OpenAI API调用失败: " + response.code());
            }
            
            String responseBody = response.body().string();
            JsonObject jsonResponse = gson.fromJson(responseBody, JsonObject.class);
            
            // 解析embeddings
            List<float[]> embeddings = new ArrayList<>();
            var dataArray = jsonResponse.getAsJsonArray("data");
            
            for (int i = 0; i < dataArray.size(); i++) {
                var embeddingObj = dataArray.get(i).getAsJsonObject();
                var embeddingArray = embeddingObj.getAsJsonArray("embedding");
                
                float[] embedding = new float[embeddingArray.size()];
                for (int j = 0; j < embeddingArray.size(); j++) {
                    embedding[j] = embeddingArray.get(j).getAsFloat();
                }
                embeddings.add(embedding);
            }
            
            return embeddings;
        }
    }
    
    /**
     * 计算余弦相似度
     */
    public double cosineSimilarity(float[] vec1, float[] vec2) {
        if (vec1.length != vec2.length) {
            throw new IllegalArgumentException("向量维度不匹配");
        }
        
        double dotProduct = 0.0;
        double norm1 = 0.0;
        double norm2 = 0.0;
        
        for (int i = 0; i < vec1.length; i++) {
            dotProduct += vec1[i] * vec2[i];
            norm1 += vec1[i] * vec1[i];
            norm2 += vec2[i] * vec2[i];
        }
        
        if (norm1 == 0.0 || norm2 == 0.0) {
            return 0.0;
        }
        
        return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2));
    }
    
    /**
     * 将float数组转换为JSON字符串
     */
    public String embeddingToJson(float[] embedding) {
        return gson.toJson(embedding);
    }
    
    /**
     * 将JSON字符串转换为float数组
     */
    public float[] jsonToEmbedding(String json) {
        return gson.fromJson(json, float[].class);
    }
    
    /**
     * 获取embedding维度
     */
    public int getDimension() {
        return dimension;
    }
    
    /**
     * 获取模型名称
     */
    public String getModel() {
        return model;
    }
}

